package com.gitee.luosl.cola.stream;

import com.gitee.luosl.cola.lang.AbstractIterator;
import com.gitee.luosl.cola.lang.GroupedIterator;
import com.gitee.luosl.cola.lang.Tuple;
import com.gitee.luosl.cola.lang.Tuple2;
import com.gitee.luosl.cola.util.CollUtils;

import java.util.*;
import java.util.function.*;
import java.util.stream.*;

public interface StreamWrapper<E> extends AutoCloseable {

    /**
     * 获取被包装的 stream
     * @return 被包装的 stream
     */
    Stream<E> getStream();

    default void forEach(Consumer<? super E> action) {
        getStream().forEach(action);
    }

    default void forEachOrdered(Consumer<? super E> action) {
        getStream().forEachOrdered(action);
    }

    default Object[] toArray() {
        return getStream().toArray();
    }

    default <A> A[] toArray(IntFunction<A[]> generator) {
        return getStream().toArray(generator);
    }

    default E reduce(E identity, BinaryOperator<E> accumulator) {
        return getStream().reduce(identity, accumulator);
    }

    default Optional<E> reduce(BinaryOperator<E> accumulator) {
        return getStream().reduce(accumulator);
    }

    default <U> U reduce(U identity, BiFunction<U, ? super E, U> accumulator, BinaryOperator<U> combiner) {
        return getStream().reduce(identity, accumulator, combiner);
    }

    default <R> R collect(Supplier<R> supplier, BiConsumer<R, ? super E> accumulator, BiConsumer<R, R> combiner) {
        return getStream().collect(supplier, accumulator, combiner);
    }

    default <R, A> R collect(Collector<? super E, A, R> collector) {
        return getStream().collect(collector);
    }

    default Optional<E> min(Comparator<? super E> comparator) {
        return getStream().min(comparator);
    }

    default Optional<E> max(Comparator<? super E> comparator) {
        return getStream().max(comparator);
    }

    default long count() {
        return getStream().count();
    }

    default boolean anyMatch(Predicate<? super E> predicate) {
        return getStream().anyMatch(predicate);
    }

    default boolean allMatch(Predicate<? super E> predicate) {
        return getStream().anyMatch(predicate);
    }

    default boolean noneMatch(Predicate<? super E> predicate) {
        return getStream().noneMatch(predicate);
    }

    default Optional<E> find() {
        return getStream().findFirst();
    }

    default Optional<E> findAny() {
        return getStream().findAny();
    }

    default AbstractIterator<E> iterator() {
        return AbstractIterator.wrap(getStream().iterator());
    }

    default boolean isParallel(){
        return getStream().isParallel();
    }

    default Stream<E> onClose(Runnable closeHandler) {
        return getStream().onClose(closeHandler);
    }

    @Override
    default void close() {
        getStream().close();
    }

    // ADD API
    default Optional<E> headOption() {
        return find();
    }

    default List<E> toList(){
        return getStream().collect(Collectors.toList());
    }

    default Set<E> toSet(){
        return getStream().collect(Collectors.toSet());
    }

    default <C extends Collection<E>> C toCollection(Supplier<C> collectionFactory){
        return getStream().collect(Collectors.toCollection(collectionFactory));
    }

    default <K, V> Map<K, V> toMap(Function<E, K> keyExtractor, Function<E, V> valueExtractor) {
        return toMap(keyExtractor, valueExtractor, (v1, v2) -> v2);
    }

    default <K, V> Map<K, V> toMap(Function<E, K> keyExtractor, Function<E, V> valueExtractor, BinaryOperator<V> mergeFunction) {
        return toMap(keyExtractor, valueExtractor, mergeFunction, HashMap::new);
    }

    default <K, V, M extends Map<K, V>> Map<K, V> toMap(Function<E, K> keyExtractor, Function<E, V> valueExtractor, BinaryOperator<V> mergeFunction, Supplier<M> mapSupplier) {
        return getStream().collect(Collectors.toMap(keyExtractor, valueExtractor, mergeFunction, mapSupplier));
    }

    /**
     * 为 stream 加上当前索引，并且 StreamWrapper 会变更为 EntryStreamWrapper<T, Integer>
     * 例如:
     * GeneralStreamWrapper<String> stream = StreamWrapper.of("a", "b", "c");
     * stream.zipWithIndex().forEach((item, idx) -> System.out.println(idx + ":" + item));
     * 输出:
     * 0:a
     * 1:b
     * 2:c
     * @return 添加索引后的 EntryStreamWrapper
     */
    default EntryStreamWrapper<E, Integer> withIndex() {
        Iterator<E> itr = iterator();
        AbstractIterator<Tuple2<E, Integer>> indexedItr;
        if(!itr.hasNext()){
            indexedItr = AbstractIterator.empty();
        }else {
            indexedItr = new AbstractIterator<Tuple2<E, Integer>>() {
                private int index = 0;

                @Override
                public boolean hasNext() {
                    return itr.hasNext();
                }

                @Override
                public Tuple2<E, Integer> getNext() {
                    return Tuple.of(itr.next(), index++);
                }
            };
        }

        return wrapEntry(indexedItr.stream());
    }

    /**
     * 取 stream 中,第一个满足条件的元素，当 stream 中没有满足条件的元素时，会返回 None
     * 例如：
     * Optional<String> r = StreamWrapper.of("3", "2", "1", "1").takeWhile("2"::equals);
     * System.out.println(r);
     * 输出: Optional[2]
     * @param predicate 判定条件
     * @return 找到的结果
     */
    default Optional<E> findFirst(Predicate<? super E> predicate){
        return getStream().filter(predicate).findFirst();
    }

    /**
     * 从左开始折叠整个 stream
     * 例如,使用 foldLeft 求和:
     * GeneralStreamWrapper<Integer> stream = StreamWrapper.of(2, 1, 3);
     * Integer r = stream.foldLeft(0, (collector, item) -> collector + item);
     * System.out.println(r);
     * 输出: 6
     * @param zero 起始元素
     * @param f 折叠的操作
     * @param <U> 折叠元素类型
     * @return 在折叠后的结果
     */
    default <U> U fold(U zero, BiFunction<U, E, U> f){
        U xs = zero;
        Iterator<E> itr = iterator();
        while(itr.hasNext()){
            xs = f.apply(xs, itr.next());
        }
        return xs;
    }

    /**
     * 分组，并对分组元素进行映射
     * @param keyExtractor 分组 key 抽取器
     * @param mapping 分组集合元素映射函数
     * @param <K> 分组 key 类型
     * @param <R> 分组后元素类型
     * @return Map<K, List<R>>
     */
    default <K, R> Map<K, List<R>> groupBy(Function<E, K> keyExtractor, Function<E, R> mapping) {
        return groupBy(keyExtractor, mapping, ArrayList::new);
    }

    default <K, R> EntryStreamWrapper<K, List<R>> groupByWrap(Function<E, K> keyExtractor, Function<E, R> mapping){
        return wrapEntry(groupBy(keyExtractor, mapping));
    }

    /**
     * 分组 聚合为指定 集合
     * @param keyExtractor 分组 key 抽取器
     * @param mapping 分组集合元素映射函数
     * @param collectionFactory  分组后集合类型提供者
     * @param <C> 结果集合类型
     * @param <K> 分组 key 类型
     * @param <R> 分组后元素类型
     * @return Map<K, C>
     */
    default <C extends Collection<R>, K, R> Map<K, C> groupBy(Function<E, K> keyExtractor,
                                                              Function<E, R> mapping,
                                                              Supplier<C> collectionFactory) {
        return getStream().collect(Collectors.groupingBy(
                keyExtractor,
                Collectors.mapping(mapping, Collectors.toCollection(collectionFactory))
        ));
    }

    default <C extends Collection<R>, K, R> EntryStreamWrapper<K, C> groupByWrap(Function<E, K> keyExtractor,
                                                                                 Function<E, R> mapping,
                                                                                 Supplier<C> collectionFactory) {
        return wrapEntry(groupBy(keyExtractor, mapping, collectionFactory));
    }


    /**
     * 对 stream 进行分组
     * 例如，按对2的求余结果进行分组:
     * GeneralStreamWrapper<Integer> stream = StreamWrapper.of(2, 1, 3, 4);
     * Map<Integer, List<Integer>> r = stream.groupBy(i -> i % 2);
     * System.out.println(r);
     * 输出: {0=[2, 4], 1=[1, 3]}
     * @param keyExtractor 分组 key 抽取器
     * @param <K> 分组 key 类型
     * @return
     */
    default <K> Map<K, List<E>> groupBy(Function<E, K> keyExtractor) {
        return getStream().collect(Collectors.groupingBy(keyExtractor));
    }

    default <K> EntryStreamWrapper<K, List<E>> groupByWrap(Function<E, K> keyExtractor) {
        return wrapEntry(groupBy(keyExtractor));
    }

    /**
     * 对 stream 进行聚合
     * 例如，分组求均值:
     * GeneralStreamWrapper<Integer> stream = StreamWrapper.of(2, 1, 3, 4);
     * Map<Integer, OptionalDouble> r = stream.agg(i -> i % 2, list -> list.stream().mapToInt(i -> i).average());
     * System.out.println(r);
     * 输出: {0=OptionalDouble[3.0], 1=OptionalDouble[2.0]}
     * @param keyExtractor 分组 key 抽取器
     * @param aggMapper 聚合函数
     * @param <K> 分组 key 类型
     * @param <V> 聚合结果 key 类型
     * @return 聚合后的结果
     */
    default <K, V> Map<K, V> agg(Function<E, K> keyExtractor, Function<List<E>, V> aggMapper){
        return aggWrap(keyExtractor, aggMapper).toMap();
    }

    default <K, V> EntryStreamWrapper<K, V> aggWrap(Function<E, K> keyExtractor, Function<List<E>, V> aggMapper){
        return wrapEntry(groupBy(keyExtractor)).mapValues(aggMapper);
    }

    /**
     * 对流构造字符串
     * 例如:
     * GeneralStreamWrapper<Integer> stream = StreamWrapper.of(2, 1, 3, 4);
     * String r = stream.mkStr(",", i -> "number:" + i, "[", "]");
     * System.out.println(r);
     * 输出: [number:2,number:1,number:3,number:4]
     * @param delimiter 分隔符
     * @param strFun 元素转字符串函数
     * @param prefix 前缀
     * @param suffix 后缀
     * @return 构造的字符串
     */
    default String mkStr(CharSequence delimiter, Function<E, String> strFun, CharSequence prefix, CharSequence suffix){
        return getStream().map(strFun).collect(Collectors.joining(delimiter, prefix, suffix));
    }

    default String mkStr(CharSequence delimiter, Function<E, String> strFun) {
        return mkStr(delimiter, strFun, "", "");
    }

    default String mkStr(CharSequence delimiter){
        return mkStr(delimiter, Objects::toString);
    }

    /**
     * 求 stream 中最大的 topN 元素
     * GeneralStreamWrapper<Integer> stream = StreamWrapper.of(2, 1, 3, 4);
     * Collection<Integer> res = stream.maxTopN(2);
     * System.out.println(res);
     * 输出: [3, 4]
     * @param topN 取多少个 top 元素
     * @param comparator 比较器
     * @return topN 集合
     */
    default Collection<E> maxTopN(int topN, Comparator<E> comparator){
       return CollUtils.maxTopN(iterator(), topN, comparator);
    }

    /**
     * 求 stream 中最小的 topN 元素
     * GeneralStreamWrapper<Integer> stream = StreamWrapper.of(2, 1, 3, 4);
     * Collection<Integer> res = stream.minTopN(2);
     * System.out.println(res);
     * 输出: [1, 2]
     * @param topN 取多少个 top 元素
     * @param comparator 比较器
     * @return topN 集合
     */
    default Collection<E> minTopN(int topN, Comparator<E> comparator){
        return maxTopN(topN, comparator.reversed());
    }

    default Collection<E> maxTopN(int topN){
        @SuppressWarnings("unchecked")
        Comparator<E> comp = (Comparator<E>) Comparator.naturalOrder();
        return maxTopN(topN, comp);
    }

    default Collection<E> minTopN(int topN){
        @SuppressWarnings("unchecked")
        Comparator<E> comp = (Comparator<E>) Comparator.naturalOrder();
        return minTopN(topN, comp);
    }


    default GeneralStreamWrapper<List<E>> sliding(int size, int step){
        return CollUtils.sliding(iterator(), size, step).streamWrapper();
    }

    default GeneralStreamWrapper<List<E>> sliding(int size){
        return sliding(size, 1);
    }

    default GeneralStreamWrapper<List<E>> partition(int size){
        return sliding(size);
    }

    default Tuple2<List<E>, List<E>> span(Predicate<? super E> predicate) {
        Map<Boolean, List<E>> grouped = groupBy(predicate::test);
        Function<Boolean, List<E>> fun = key -> grouped.getOrDefault(key, Collections.emptyList());
        return Tuple.of(fun.apply(true), fun.apply(false));
    }

    default Tuple2<List<E>, List<E>> splitAt(int n) {
        List<E> left = new ArrayList<>();
        List<E> right = new ArrayList<>();
        AbstractIterator<E> itr = iterator();
        int current = 0;
        while (current < n && itr.hasNext()){
            left.add(itr.next());
            current++;
        }
        while(itr.hasNext()){
            right.add(itr.next());
        }
        return Tuple.of(left, right);
    }

    default <C extends Collection<E>> C take(int n, Supplier<C> supplier) {
        return wrap(getStream().limit(n)).toCollection(supplier);
    }

    default List<E> take(int n) {
        return take(n, ArrayList::new);
    }

    /*** 工厂方法 ***/

    @SafeVarargs
    static <T> GeneralStreamWrapper<T> of(T ...values){
        return wrap(Stream.of(values));
    }

    @SafeVarargs
    static <K, V> EntryStreamWrapper<K, V> ofEntry(Tuple2<K, V> ...values){
        return wrapEntry(Stream.of(values));
    }

    static <T> GeneralStreamWrapper<T> wrap(Stream<T> stream){
        return new GeneralStreamWrapper<>(stream);
    }

    static <T> GeneralStreamWrapper<T> wrap(Collection<T> collection){
        return wrap(collection.stream());
    }

    static <T> GeneralStreamWrapper<T> wrap(Iterator<T> itr){
        return wrap(AbstractIterator.wrap(itr).stream());
    }

    static <T> GeneralStreamWrapper<T> wrap(final T[] array) {
        return wrap(Arrays.stream(array));
    }

    static <E, K, V> EntryStreamWrapper<K, V> wrapEntry(Collection<E> coll, Function<E, Tuple2<K, V>> transformer){
        return wrapEntry(coll.stream(), transformer);
    }

    static <E, K, V> EntryStreamWrapper<K, V> wrapEntry(Collection<E> coll, Function<E, K> keyMapper, Function<E, V> valueMapper){
        return wrapEntry(coll, item -> Tuple.of(keyMapper.apply(item), valueMapper.apply(item)));
    }

    static <K, V> EntryStreamWrapper<K, V> wrapEntry(Collection<Tuple2<K, V>> coll){
        return new EntryStreamWrapper<>(coll.stream());
    }

    static <K, V> EntryStreamWrapper<K, V> wrapEntry(Stream<Tuple2<K, V>> stream){
        return new EntryStreamWrapper<>(stream);
    }

    static <T, K, V> EntryStreamWrapper<K, V> wrapEntry(Stream<T> stream, Function<T, Tuple2<K, V>> transformer){
        return wrapEntry(stream.map(transformer));
    }

    static <K, V> EntryStreamWrapper<K, V> wrapEntry(Map<K, V> map){
        return wrapEntry(map.entrySet().stream(), Tuple::of);
    }

    static <K, V> EntryStreamWrapper<K, V> concat(EntryStreamWrapper<K, V> left, EntryStreamWrapper<K, V> right) {
        return wrapEntry(Stream.concat(left.getStream(), right.getStream()));
    }

    static <T> GeneralStreamWrapper<T> concat(GeneralStreamWrapper<T> left, GeneralStreamWrapper<T> right) {
        return wrap(Stream.concat(left.getStream(), right.getStream()));
    }

}
